import os

from typing import Optional, List
from fastapi import FastAPI, Form
from pydantic import BaseModel

import base64
import cv2
import requests
# from face_embedding.register_pkl import register_pkl_embedding, re_register_pkl_embedding
# from face_embedding.one2one_test import one2onne_embedding
from face_embedding.face_crop import FaceRetinaModel
from face_embedding.utils import argsutils
from face_embedding.face_feature_thread import FaceFeatureThread
from utils import dictProperties, dbpool_util, log_util

import numpy as np

import traceback
from datetime import datetime
import json
import re
import time
import random
import uuid
import math

from face_embedding.face_feature_tvm import TVMGetFeatureModel

app = FastAPI()

faiss_server_url = dictProperties.get('faiss_server')

args = argsutils.Facemodel(dictProperties).args
retina_model = FaceRetinaModel(dictProperties)
face_feature = FaceFeatureThread(args, dictProperties)
tvm_feature = TVMGetFeatureModel()

def get_db():
    return dbpool_util

def query_db(query, args=(), one=False):
    # print('sql_s1', query, file=sys.stdout)
    result = get_db().execute_query(query, args)
    return result

def execute_db(sql, args=()):
    return get_db().execute_iud(sql, args)



#####################
def insert_home_info(group_id):
    log_util.logger.debug('添加组信息开始')
    insert_sql = """insert into t_home (group_id, group_remark, create_time)values ('{}','{}', now());""".format(group_id, '')
    count = execute_db(insert_sql)
    log_util.logger.debug('添加组信息结束')

def delete_home_info(group_id):
    log_util.logger.debug('删除组信息开始')
    delete_sql = """delete from t_home where group_id='{}'""".format(group_id)
    count = execute_db(delete_sql)
    log_util.logger.debug('删除组信息结束')

def get_home_info(group_id):
    log_util.logger.debug('查询组信息开始')
    sql_s = "select home_id, group_id, group_remark, create_time from t_home where group_id='{group_id}';".format(
        group_id=group_id)
    group_info = query_db(sql_s)
    log_util.logger.debug('查询组信息结束')
    if group_info and len(group_info) > 0:
        return {
            'home_id': group_info[0][0],
            'group_id': group_info[0][1],
            'group_remark': group_info[0][2],
            'create_time': group_info[0][3]
        }
    else:
        return None


def insert_face(group_id, feature_data, ext_id):
    log_util.logger.debug('添加人脸注册信息开始')
    embedding_json = json.dumps(feature_data)
    face_id = str(uuid.uuid1())
    insert_sql = """
                    insert into t_home_user_face (group_id, face_id, feature_data, ext_id, create_time)
                    values ('{}','{}','{}','{}', now());
                    """.format(group_id, face_id, embedding_json, ext_id)
    count = execute_db(insert_sql)
    log_util.logger.debug('添加人脸注册信息结束')
    return face_id, count

def search_face_info_by_face_id(face_id):
    log_util.logger.debug('查询人脸注册信息开始')
    sql_s = """
        select id, group_id, face_id, feature_data, ext_id, create_time from t_home_user_face 
        where face_id='{face_id}';""".format(face_id=face_id)
    face_info = query_db(sql_s)
    log_util.logger.debug('查询人脸注册信息结束')
    if face_info and len(face_info) > 0:
        return {
            'id': face_info[0][0],
            'group_id': face_info[0][1],
            'face_id': face_info[0][2],
            'feature_data': json.loads(face_info[0][3]),
            'ext_id': face_info[0][4],
            'create_time': face_info[0][5]
        }
    else:
        return None

def search_face_info(group_id, face_id, ext_id = None):
    log_util.logger.debug('查询人脸注册信息开始')
    sql_s = """
        select id, group_id, face_id, feature_data, ext_id, create_time from t_home_user_face 
        where group_id='{group_id}' and face_id='{face_id}';""".format(
        group_id=group_id, face_id=face_id)
    face_info = query_db(sql_s)
    log_util.logger.debug('查询人脸注册信息结束')
    if face_info and len(face_info) > 0:
        return {
            'id': face_info[0][0],
            'group_id': face_info[0][1],
            'face_id': face_info[0][2],
            'feature_data': json.loads(face_info[0][3]),
            'ext_id': face_info[0][4],
            'create_time': face_info[0][5]
        }
    else:
        return None

def delete_face_info(group_id, face_id, ext_id = None):
    log_util.logger.debug('删除人脸注册信息开始')
    delete_sql = """delete from t_home_user_face where group_id='{}' and face_id='{}'""".format(group_id, face_id)
    count = execute_db(delete_sql)
    log_util.logger.debug('删除人脸注册信息结束')

def insert_face_feature_log(feature_data, file_path):
    log_util.logger.debug('人脸特征提取-插入特征日志开始')
    embedding_json = json.dumps(feature_data)
    insert_sql = """ 
     insert into t_home_face_feature_log (feature_data, file_path, create_time)
                    values ('{}','{}', now());
     """.format(embedding_json, file_path)
    count = execute_db(insert_sql)
    log_util.logger.debug('人脸特征提取-插入特征日志结束')

def post_server(post_url, data):
    log_util.logger.debug('faiss服务' + post_url + '开始')
    res = requests.post(url=faiss_server_url + post_url, json=data)
    res_val = json.loads(res.text)
    log_util.logger.debug('faiss服务' + post_url + '结束')
    if res_val.get('code') == 1:
        return None, res_val.get('data')
    else:
        return ('500', 'faiss系统异常'), None

def img_resize(img):
    h, w, c = img.shape
    if h > 300 or w > 300:
        h_scale = 300
        w_scale = int(h_scale / h * w)
        img = cv2.resize(img, (w_scale, h_scale))
    return img

def create_success(message, data = None):
    if data:
        pass
    return {
                'code' : '10000'
            }

def create_error(error, message):
    return {
                'code' : error[0],
                'message' : error[1]
            }


class FaceImageParam(BaseModel):
    image_base64: str


# 人脸特征提取
@app.post("/get_face_feature")
async def get_face_feature(face_param: FaceImageParam):
    try:
        error_message_type = '人脸特征提取错误'
        log_util.logger.debug('人脸特征提取-接收到请求')
        image_base64 = face_param.image_base64
        if image_base64 is None:
            return create_error(('400', '图片信息为空'), error_message_type)

        base64Image = image_base64.replace("%2B", "+").replace("%3D", "=").replace("%2F", "/")
        img = base64.b64decode(base64Image)
        img_array = np.frombuffer(img, np.uint8)
        ori_img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
        ori_img = img_resize(ori_img)
        log_util.logger.debug('人脸特征提取-图片预处理完成')
        dir = './imgs/get_feature'  # + '/' + userCode
        if not os.path.exists(dir):
            os.makedirs(dir)
        num = datetime.now().strftime("%m-%d-%Y-%H-%M-%S") + '-' + str(random.randint(1, 10))
        # 写原图
        img_path = dir + '/' + num + ".jpg"
        cv2.imwrite(img_path, ori_img)
        log_util.logger.debug('人脸特征提取-图片保存完成')
        error, face_imgs, bboxs, is_masks, face_angles = retina_model.detect_face(ori_img)
        if error:
            return create_error(error, error_message_type)

        faces_result = []
        for (i, face_img) in enumerate(face_imgs):
            image = face_img
            # convert face to RGB color
            nimg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            nimg = cv2.resize(nimg, (112, 112))
            nimg = np.transpose(nimg, (2, 0, 1))
            error, face_embedding = face_feature.get_feature(nimg)
            if error:
                return create_error(error, error_message_type)
            faces_result.append({
                'bounding_box': bboxs[i],
                'feature_Data': face_embedding.tolist(),
                'score': 0.9,
                'face_num': len(face_imgs),
                'face_angle': face_angles[i],
                'is_masked': is_masks[i]
            })
            insert_face_feature_log(face_embedding.tolist(), img_path)
        result_data = {
            'code': '10000',
            'faces': faces_result
        }
        log_util.logger.debug('人脸特征提取-接收到请求完成')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

# 人脸特征提取TVM
@app.post("/get_face_feature_tvm")
async def get_face_feature_tvm(face_param: FaceImageParam):
    try:
        error_message_type = '人脸特征提取错误'
        log_util.logger.debug('人脸特征提取-接收到请求')
        image_base64 = face_param.image_base64
        if image_base64 is None:
            return create_error(('400', '图片信息为空'), error_message_type)

        base64Image = image_base64.replace("%2B", "+").replace("%3D", "=").replace("%2F", "/")
        img = base64.b64decode(base64Image)
        img_array = np.frombuffer(img, np.uint8)
        ori_img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
        ori_img = img_resize(ori_img)
        log_util.logger.debug('人脸特征提取-图片预处理完成')
        dir = './imgs/get_feature'  # + '/' + userCode
        if not os.path.exists(dir):
            os.makedirs(dir)
        num = datetime.now().strftime("%m-%d-%Y-%H-%M-%S") + '-' + str(random.randint(1, 10))
        # 写原图
        img_path = dir + '/' + num + ".jpg"
        cv2.imwrite(img_path, ori_img)
        log_util.logger.debug('人脸特征提取-图片保存完成')
        error, face_imgs, bboxs, is_masks, face_angles = retina_model.detect_face(ori_img)
        if error:
            return create_error(error, error_message_type)

        faces_result = []
        for (i, face_img) in enumerate(face_imgs):
            image = face_img
            # convert face to RGB color
            nimg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            nimg = cv2.resize(nimg, (112, 112))
            nimg = np.transpose(nimg, (2, 0, 1))
            error, face_embedding = tvm_feature.get_feature(nimg)
            if error:
                return create_error(error, error_message_type)
            faces_result.append({
                'bounding_box': bboxs[i],
                'feature_Data': face_embedding.tolist(),
                'score': 0.9,
                'face_num': len(face_imgs),
                'face_angle': face_angles[i],
                'is_masked': is_masks[i]
            })
            insert_face_feature_log(face_embedding.tolist(), img_path)
        result_data = {
            'code': '10000',
            'faces': faces_result
        }
        log_util.logger.debug('人脸特征提取-接收到请求完成')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

class FaceSearchNParam(BaseModel):
    feature_data: List[float]
    group_id: str
    top_n: int = 1
    threshold: float = 0

# 人脸搜索（1：N）
@app.post("/get_search_n")
async def get_search_n(face_param: FaceSearchNParam):
    try:
        log_util.logger.debug('人脸搜索(1:N)-接收到请求')
        error_message_type = '人脸搜索（1：N）错误'
        feature_data = face_param.feature_data
        group_id = face_param.group_id
        top_n = face_param.top_n
        threshold = face_param.threshold

        group_info = get_home_info(group_id)
        if group_info is None:
            return create_error(('400', '人脸库不存在'), error_message_type)

        url_str = 'face_search'
        data = {'groupId': group_id, 'top_n': top_n, 'threshold': threshold, 'embedding': json.dumps(feature_data)}

        error, result_data = post_server(url_str, data)
        if error:
            return create_error(error, error_message_type)
        print(result_data)
        faces_result = []
        if result_data:
            face_ids = result_data.get('userCode')
            I = result_data.get('I')
            D = result_data.get('D')
            # print('I:', I)
            print('D:', D)
            if face_ids:
                index = 0
                for face_id in face_ids:
                    user_info = search_face_info_by_face_id(face_id)
                    faces_result.append({
                        'face_id': face_id,
                        'ext_id': user_info.get('ext_id'),
                        'score': (2.0 - D[0][index]) / 2,
                        # 'score': 1.0 - (math.sqrt(D[0][index]) / 2)
                        # 'score': 1.0 / (1.0 + D[0][index])
                    })
                    index += 1
        result_data = {
            'code': '10000',
            'faces': faces_result
        }
        log_util.logger.debug('人脸搜索(1:N)-接收到请求处理完成')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

# 创建人脸库
class FaceDBParam(BaseModel):
    group_id: str

@app.post("/create_facedb")
async def create_facedb(face_param: FaceDBParam):
    try:
        log_util.logger.debug('创建人脸库-接收到请求')
        error_message_type = '创建人脸库错误'
        group_id = face_param.group_id

        group_info = get_home_info(group_id)
        if group_info:
            return create_error(('400', '人脸库已存在'), error_message_type)
        insert_home_info(group_id)
        result_data = {
            'code': '10000'
        }
        log_util.logger.debug('创建人脸库-接收到请求处理完成')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

# 查询人脸库
@app.post("/search_facedb")
async def search_facedb(face_param: FaceDBParam):
    try:
        log_util.logger.debug('查询人脸库-接收到请求')
        error_message_type = '查询人脸库错误'
        group_id = face_param.group_id

        group_info = get_home_info(group_id)
        if group_info:
            result_data = {
                'code': '10000',
                'group_id': group_info.get('group_id'),
                'create_time': group_info.get('create_time')
            }
        else:
            result_data = {
                'code': '10000'
            }
        log_util.logger.debug('查询人脸库-接收到请求处理完成')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

# 删除人脸库
@app.post("/delete_facedb")
async def delete_facedb(face_param: FaceDBParam):
    try:
        log_util.logger.debug('删除人脸库-接收到请求')
        error_message_type = '删除人脸库错误'
        group_id = face_param.group_id

        delete_home_info(group_id)

        url_str = 'delete_group'
        data = {'groupId': group_id}

        error, result_data = post_server(url_str, data)
        log_util.logger.debug('删除人脸库-接收到请求处理结束')
        if error:
            return create_error(('500', '系统异常'), error_message_type)

        result_data = {
            'code': '10000'
        }
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

# 人脸特征管理

class FaceFeatureAddParam(BaseModel):
    group_id: str
    feature_data: List[float]
    ext_id: str = None

# 添加人脸
@app.post("/create_face")
async def create_face(face_param: FaceFeatureAddParam):
    try:
        log_util.logger.debug('添加人脸-接收到请求')
        error_message_type = '添加人脸错误'
        group_id = face_param.group_id
        feature_data = face_param.feature_data
        ext_id = face_param.ext_id

        group_info = get_home_info(group_id)
        if group_info is None:
            return create_error(('400', '人脸库不存在'), error_message_type)

        if not ext_id:
            ext_id = uuid.uuid1()

        face_id, count = insert_face(group_id, feature_data, ext_id)

        url_str = 'face_add'
        data = {'groupId': group_id, 'faceId': face_id, 'embedding': json.dumps([feature_data])}

        error, result_data = post_server(url_str, data)
        if error:
            delete_face_info(group_id, face_id, ext_id)
            return create_error(('500', '系统异常'), error_message_type)

        result_data = {
            'code': '10000',
            'face_id': face_id
        }
        log_util.logger.debug('添加人脸-接收到请求处理完成')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)
# 查询人脸

class FaceFeatureSearchParam(BaseModel):
    group_id: str
    face_id: str
    ext_id: str = None

@app.post("/search_face")
async def search_face(face_param: FaceFeatureSearchParam):
    try:
        log_util.logger.debug('查询人脸-接收到请求')
        error_message_type = '查询人脸错误'
        group_id = face_param.group_id
        face_id = face_param.face_id
        ext_id = face_param.ext_id

        face_info = search_face_info(group_id, face_id, ext_id)
        feature_data = None
        if face_info:
            feature_data = face_info.get('feature_data')

        result_data = {
            'code': '10000',
            'face_id': face_id,
            'ext_id': ext_id,
            'feature_data': feature_data
        }
        log_util.logger.debug('查询人脸-接收到请求处理结束')
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)
# 删除人脸
@app.post("/delete_face")
async def delete_face(face_param: FaceFeatureSearchParam):
    try:
        log_util.logger.debug('删除人脸-接收到请求')
        error_message_type = '删除人脸错误'
        group_id = face_param.group_id
        face_id = face_param.face_id
        ext_id = face_param.ext_id

        face_info = search_face_info(group_id, face_id, ext_id)
        if face_info is None:
            return create_error(('400', '找不到对应数据'), error_message_type)

        delete_face_info(group_id, face_id, ext_id)
        url_str = 'face_reload'
        data = {'groupId': group_id}
        error, result_data = post_server(url_str, data)
        log_util.logger.debug('删除人脸-接收到请求处理完成')
        if error:
            return create_error(error, '更新失败')

        result_data = {
            'code': '10000'
        }
        return result_data

    except Exception as e:
        traceback.print_exc()  ###
        return create_error(('500', '系统异常'), error_message_type)

####------------------------------------------------------------------------------------
@app.get("/")
def index():
    log_util.logger.debug('首页-接收到请求')
    return {"message": "oaclockin service is runing"}



# export LC_ALL=en_US.utf-8
# export LANG=en_US.utf-8

# uvicorn api_oa_face_fastapi_v3:app --reload --host 0.0.0.0 --port 17104

# uvicorn api_oa_face_fastapi_v3:app --host 0.0.0.0 --port 17104 --workers 4